import os
import random
import argparse
import shutil
from shutil import copy2
import urllib
import traceback
import time
import sys

SRC_PATH = ''
DST_PATH = ''
TRAIN = 7
TEST = 0
WW_MAX = 0

def allImgFilePath(rootPath, allFIleList):
    fileList = os.listdir(rootPath)
    for temp in fileList:
        if os.path.isfile(os.path.join(rootPath,temp)):
            if temp.endswith(".jpg"):
                allFIleList.append(os.path.join(rootPath,temp))
        else:
            allImgFilePath(os.path.join(rootPath,temp),allFIleList)


def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)
        print('mkdir path: {}'.format(path))      


def rand_mov_write_file(fileImgList, dstPath):
    count = 0
    if len(fileImgList) > 0:
        train_file = '{}/train.txt'.format(dstPath)
        val_file = '{}/val.txt'.format(dstPath)
        imgs_path = '{}/images'.format(dstPath)
        labels_path = '{}/labels'.format(dstPath)        
        
        train_f = open(train_file, 'w')
        val_f = open(val_file, 'w')
        train_count = 0
        val_count = 0
        test_count = 0
        err_count = 0
        for img in fileImgList:
            #print('{} img:{}'.format(count+1, img))
            txt= img.replace(".jpg",".txt")
            if not os.path.exists(txt):
                err_count = err_count + 1
                print('{} txt not exitst:{}'.format(err_count, txt))
                continue
            
            if count%10<TRAIN and not img.endswith("val_ww.jpg"):
                dst_img_path = '{}/t_{}.jpg'.format(imgs_path, train_count+1)
                dst_txt_path = '{}/t_{}.txt'.format(labels_path, train_count+1)
                shutil.copy2(img, dst_img_path)
                shutil.copy2(txt, dst_txt_path)
                train_count = train_count + 1
                train_f.write(dst_img_path + "\n")
            else:
                dst_img_path = '{}/v_{}.jpg'.format(imgs_path, val_count+1)
                dst_txt_path = '{}/v_{}.txt'.format(labels_path, val_count+1)
                shutil.copy2(img, dst_img_path)
                shutil.copy2(txt, dst_txt_path)
                val_count = val_count + 1
                val_f.write(dst_img_path + "\n")

            count = count + 1

            if TEST>0 and count%TEST==0:
                dst_img_path = '{}/test_{}.jpg'.format(imgs_path, test_count+1)
                shutil.copy2(img, dst_img_path)
                test_count = test_count + 1
                 
            if WW_MAX>0 and count>=WW_MAX:
                break;

        train_f.close()   
        val_f.close() 
        print('count:{} train:{} val:{} test:{} err:{}'.format(count, train_count, val_count, test_count, err_count))
    return count


'''
def allTxtFilePath(rootPath, allFIleList):
    fileList = os.listdir(rootPath)
    for temp in fileList:
        if os.path.isfile(os.path.join(rootPath,temp)):
            if temp.endswith(".txt"):
                allFIleList.append(os.path.join(rootPath,temp))
        else:
            allTxtFilePath(os.path.join(rootPath,temp),allFIleList)

def allFilePath(rootPath, allImgFIleList, allTxtFileList):
    fileList = os.listdir(rootPath)
    for temp in fileList:
        if os.path.isfile(os.path.join(rootPath,temp)):
            if temp.endswith(".jpg"):
                allImgFIleList.append(os.path.join(rootPath,temp))
            if temp.endswith(".txt"):
                allTxtFileList.append(os.path.join(rootPath,temp))                
        else:
            allFilePath(os.path.join(rootPath,temp),allImgFIleList,allTxtFileList)                        
'''

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--src', type=str, default='/root', dest='src', help='src path(s)')
    parser.add_argument('--dst', type=str, default='./', dest='dst', help='dst path(s)')
    parser.add_argument('--train', type=int, default=7, dest='train', help='')  
    parser.add_argument('--test', type=int, default=0, dest='test', help='') 
    parser.add_argument('--max', type=int, default=0, dest='max', help='')  

    opt = parser.parse_args()
    return opt


def check_arg(args):
    global SRC_PATH
    global DST_PATH
    global TRAIN
    global TEST
    global WW_MAX
    
    SRC_PATH = args.src
    print('src path: {}'.format(SRC_PATH))
    if not os.path.exists(SRC_PATH):
        print('src path not exist')
        exit(-1)

    TRAIN = args.train   
    print('TRAIN: {}'.format(TRAIN))
    TEST = args.test   
    print('TEST: {}'.format(TEST))    
    WW_MAX = args.max   
    print('max: {}'.format(WW_MAX))    

    DST_PATH = args.dst
    print('dst path: {}'.format(DST_PATH))

    mkdir(DST_PATH)
    
    imgs_path = '{}/images'.format(DST_PATH)
    mkdir(imgs_path)                
    labels_path = '{}/labels'.format(DST_PATH)
    mkdir(labels_path)           

if __name__ == '__main__':
    args = parse_opt()
    # check args 
    check_arg(args)

    img_list=[]
    allImgFilePath(SRC_PATH, img_list)
    print('file count={}'.format(len(img_list)))

    rand_mov_write_file(img_list, DST_PATH)